package com.cty.twentyEighthDay;

import java.util.*;

/*
 *给定一个二叉搜索树的根节点 root ，和一个整数 k ，请你设计一个算法查找其中第 k 小的元素（从 1 开始计数）。
 * */
public class Q_44 {
    List<Integer> ans = new ArrayList<>();

    // 中序遍历搜索 时间复杂度O(n) 空间复杂度O(n)
    public int kthSmallest(TreeNode root, int k) {
        midSort(root);
        return ans.get(k - 1);
    }

    public void midSort(TreeNode root) {
        if (root.left != null) {
            midSort(root.left);
        }
        ans.add(root.val);
        if (root.right != null) {
            midSort(root.right);
        }
    }

    // 力扣官方题解 方法一：中序遍历
    /*
     * 时间复杂度O(H + k) 其中H是树的高度
     * 空间复杂度O(N)
     *
     * */
    public int kthSmallest2(TreeNode root, int k) {
        Stack<TreeNode> stack = new Stack<>();
        while (!stack.isEmpty() || root != null) {
            while (root != null) {
                stack.push(root);
                root = root.left;
            }
            root = stack.pop();
            --k;
            if (k == 0) {
                break;
            }
            root = root.right;
        }
        return root.val;
    }

    // 力扣官方题解 方法二:记录子树节点树
    /*
     * 思路:
     *   首先我们记录根节点的左子树的节点个数 若大于 k - 1 则表示 目标节点在左子树中
     *   若大于 k - 1 则表示目标节点在右子树中 若等于 则表示当前节点就是目标节点
     *   时间复杂度O(n) 空间复杂度O(n)
     * */
    public int kthSmallest3(TreeNode root, int k) {
        MyBST myBST = new MyBST(root);
        return myBST.kthSmallest(k);
    }

    // 方法三:平衡二叉搜索树 主包看这个代码太多了 所以主包跳过了
//    public int kthSmallest4(TreeNode root, int k) {
//
//    }
}

class MyBST {
    // 根节点
    TreeNode root;
    // 记录当前树的总个数
    Map<TreeNode, Integer> map;

    public MyBST(TreeNode root) {
        this.root = root;
        map = new HashMap<>();
        // 记录节点的数量
        recordNodeNum(root);
    }

    private int recordNodeNum(TreeNode root) {
        if (root == null) {
            return 0;
        }
        // 记录左子树的数量 记录右子树的数量 二者加一 就是以当前为根节点二叉树的节点总数
        map.put(root, recordNodeNum(root.left) + recordNodeNum(root.right) + 1);
        // 返回当前节点的总数
        return map.get(root);
    }


    public int kthSmallest(int k) {
        TreeNode node = root;
        while (node != null) {
            // 获得左子树的数量
            Integer left= map.getOrDefault(node.left, 0);
            if (left > k - 1) {
                node = node.left;
            } else if(left < k - 1){
                k -= (left + 1);
                node = node.right;
            } else {
                break;
            }
        }
        return node.val;
    }
}
